#! /usr/bin/env python
# -*- coding: utf-8 -*-
# vim:fenc=utf-8
#
# Copyright © 2019 crane <crane@gosun>
#
# Distributed under terms of the MIT license.

"""
https://www.lintcode.com/problem/k-sum/description

思路来源于整数拆分: 1: 把数排序, 2: 个数k开始参数
"""

class Solution:
    """
    @param A: An integer array
    @param k: A positive integer (k <= length(A))
    @param target: An integer
    @return: An integer
    """
    def kSum(self, nums, k, target):
        self.nums = sorted(nums)
        self.len = len(self.nums)
        self.k = k
        self.target = target
        self.cnt_memory = {}
        return self.cnt_from_min(target, k, 0)

    def cnt_from_min(self, target, k, from_min_idx):
        # 最小值从from_min_idx开始匹配, 匹配target, 还剩下k个数可用
        key = (target, k, from_min_idx)
        cnt = self.cnt_memory.get(key)
        if cnt is not None:
            return cnt

        # print('key', key)

        if target == 0:
            cnt = 1
            return cnt

        if k <= 0:
            return 0

        else:
            cnt = 0
            for min_idx in range(from_min_idx, self.len-k+1):
                # print('k', k, from_min_idx, min_idx)

                remain = target - self.nums[min_idx]

                if remain < self.nums[from_min_idx] and k-1 > 0:
                    # 用来剪枝
                    # print('hit')
                    break

                # 注意这里min_idx + 1, 表示不使用重复数字
                cnt += self.cnt_from_min(remain, k-1, min_idx+1)

        self.cnt_memory[key] = cnt
        # print(from_min_idx, k, target, cnt)
        # print(key, cnt)
        return cnt


def test():
    s = Solution()

    lst = [1, 2, 3, 4]
    assert s.kSum(lst, 2, 5) == 2

    lst = [1]
    assert s.kSum(lst, 1, 1) == 1

    lst = [1, 2, 3, 4, 5]
    assert s.kSum(lst, 3, 6) == 1


def main():
    print("start main")
    test()

if __name__ == "__main__":
    main()
